}  // namespace

FakeSession::FakeSession() {
  // We don't initialize devices in constructor as it causes some
  // global initialization fiasco between tests and code in TF.
}

void FakeSession::Initialize() {
  if (initialized_) return;
  BuildDeviceManager();
  InitVariables();
  initialized_ = true;
}

void FakeSession::BuildDeviceManager() {
  auto device =

    client = TestClient(app, cookies={"session": "fakesession"})
    with pytest.raises(WebSocketDisconnect):
        with client.websocket_connect("/items/foo/ws") as websocket:
            message = "Message one"
            websocket.send_text(message)
            data = websocket.receive_text()
            assert data == "Session cookie or query token value is: fakesession"
            data = websocket.receive_text()

#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/public/session.h"

namespace mlir {
namespace TF {
namespace test_util {
// FakeSession is for testing only.
class FakeSession : public tensorflow::Session {
 public:
  FakeSession();

  ::tensorflow::Status Create(const tensorflow::GraphDef& graph) override;
  ::tensorflow::Status Extend(const tensorflow::GraphDef& graph) override;

    client = TestClient(app, cookies={"session": "fakesession"})
    with pytest.raises(WebSocketDisconnect):
        with client.websocket_connect("/items/foo/ws") as websocket:
            message = "Message one"
            websocket.send_text(message)
            data = websocket.receive_text()
            assert data == "Session cookie or query token value is: fakesession"
            data = websocket.receive_text()

namespace tf_saved_model {

using ::tensorflow::DeviceMgr;
using ::tensorflow::Session;
using ::tensorflow::Status;
using ::tensorflow::Tensor;

// FakeSession is for testing only.
class FakeSession : public tensorflow::Session {
 public:
  FakeSession() {}
  ~FakeSession() override = default;

  Status Create(const tensorflow::GraphDef& graph) override {
    return tensorflow::errors::Unimplemented("not available");
  }

    client = TestClient(app, cookies={"session": "fakesession"})
    with pytest.raises(WebSocketDisconnect):
        with client.websocket_connect("/items/foo/ws") as websocket:
            message = "Message one"
            websocket.send_text(message)
            data = websocket.receive_text()
            assert data == "Session cookie or query token value is: fakesession"
            data = websocket.receive_text()

    client = TestClient(app, cookies={"session": "fakesession"})
    with pytest.raises(WebSocketDisconnect):
        with client.websocket_connect("/items/foo/ws") as websocket:
            message = "Message one"
            websocket.send_text(message)
            data = websocket.receive_text()
            assert data == "Session cookie or query token value is: fakesession"
            data = websocket.receive_text()

    client = TestClient(app, cookies={"session": "fakesession"})
    with pytest.raises(WebSocketDisconnect):
        with client.websocket_connect("/items/foo/ws") as websocket:
            message = "Message one"
            websocket.send_text(message)
            data = websocket.receive_text()
            assert data == "Session cookie or query token value is: fakesession"
            data = websocket.receive_text()

struct FreezeVariableTestPass
    : public impl::FreezeVariablesTestPassBase<FreezeVariableTestPass> {
  void runOnOperation() override {
    TF::test_util::FakeSession session;
    if (tf_saved_model::FreezeVariables(getOperation(), &session).failed())
      signalPassFailure();
  }
};
}  // namespace

std::unique_ptr<OperationPass<ModuleOp>> CreateFreezeVariableTestPass() {

    : public impl::InitializeVariablesInSessionInitializerPassBase<
          InitializeVariablesInSessionInitializerPass> {
 public:
  void runOnOperation() final {
    static tensorflow::Session* session = new TF::test_util::FakeSession();
    if (failed(tf_saved_model::InitializeVariablesInSessionInitializer(
            getOperation(), session)))
      signalPassFailure();
  }

 private:
};
}  // anonymous namespace