4 files changed, 43 insertions, 9 deletions

diff --git a/src/tests/cmocka/common_mock_resp.c b/src/tests/cmocka/common_mock_resp.c
index ce73d1b45..dc03d39b6 100644
--- a/src/tests/cmocka/common_mock_resp.c
+++ b/src/tests/cmocka/common_mock_resp.c
@@ -63,12 +63,23 @@ mock_cctx(TALLOC_CTX *mem_ctx, struct resp_ctx *rctx)
     cctx = talloc_zero(mem_ctx, struct cli_ctx);
     if (!cctx) return NULL;
 
-    cctx->creq = talloc_zero(cctx, struct cli_request);
-    if (cctx->creq == NULL) {
-        talloc_free(cctx);
+    cctx->rctx = rctx;
+    return cctx;
+}
+
+struct cli_protocol *
+mock_prctx(TALLOC_CTX *mem_ctx)
+{
+    struct cli_protocol *prctx;
+
+    prctx = talloc_zero(mem_ctx, struct cli_protocol);
+    if (!prctx) return NULL;
+
+    prctx->creq = talloc_zero(prctx, struct cli_request);
+    if (prctx->creq == NULL) {
+        talloc_free(prctx);
         return NULL;
     }
 
-    cctx->rctx = rctx;
-    return cctx;
+    return prctx;
 }
diff --git a/src/tests/cmocka/common_mock_resp.h b/src/tests/cmocka/common_mock_resp.h
index a4d8f55c7..aab6a94e4 100644
--- a/src/tests/cmocka/common_mock_resp.h
+++ b/src/tests/cmocka/common_mock_resp.h
@@ -38,6 +38,9 @@ mock_rctx(TALLOC_CTX *mem_ctx,
 struct cli_ctx *
 mock_cctx(TALLOC_CTX *mem_ctx, struct resp_ctx *rctx);
 
+struct cli_protocol *
+mock_prctx(TALLOC_CTX *mem_ctx);
+
 /* When mocking a module that calls sss_dp_get_account_{send,recv}
  * requests, your test, when linked against this module, will call
  * the mock functions instead. Then you can simulate results of the
diff --git a/src/tests/cmocka/test_nss_srv.c b/src/tests/cmocka/test_nss_srv.c
index d0b1e28e0..945e2b0c2 100644
--- a/src/tests/cmocka/test_nss_srv.c
+++ b/src/tests/cmocka/test_nss_srv.c
@@ -119,13 +119,17 @@ static void set_cmd_cb(cmd_cb_fn_t fn)
 
 void __wrap_sss_cmd_done(struct cli_ctx *cctx, void *freectx)
 {
-    struct sss_packet *packet = cctx->creq->out;
+    struct cli_protocol *pctx;
+    struct sss_packet *packet;
     uint8_t *body;
     size_t blen;
     cmd_cb_fn_t check_cb;
 
     check_cb = sss_mock_ptr_type(cmd_cb_fn_t);
 
+    pctx = talloc_get_type(cctx->protocol_ctx, struct cli_protocol);
+    packet = pctx->creq->out;
+
     __real_sss_packet_get_body(packet, &body, &blen);
 
     nss_test_ctx->tctx->error = check_cb(sss_packet_get_status(packet),
@@ -1070,6 +1074,15 @@ void test_nss_setup(struct sss_test_conf_param params[],
     /* Create client context */
     nss_test_ctx->cctx = mock_cctx(nss_test_ctx, nss_test_ctx->rctx);
     assert_non_null(nss_test_ctx->cctx);
+
+    /* Add nss specific state_ctx */
+    nss_connection_setup(nss_test_ctx->cctx);
+    assert_non_null(nss_test_ctx->cctx->state_ctx);
+
+    /* do after previous setup as the former nulls procotol_ctx */
+    nss_test_ctx->cctx->protocol_ctx = mock_prctx(nss_test_ctx->cctx);
+    assert_non_null(nss_test_ctx->cctx->protocol_ctx);
+
 }
 
 static int test_nss_getgrnam_check(struct group *expected, struct group *gr, const int nmem)
diff --git a/src/tests/cmocka/test_pam_srv.c b/src/tests/cmocka/test_pam_srv.c
index e4ad5b650..6f56071f8 100644
--- a/src/tests/cmocka/test_pam_srv.c
+++ b/src/tests/cmocka/test_pam_srv.c
@@ -217,6 +217,7 @@ void test_pam_setup(struct sss_test_conf_param dom_params[],
                     struct sss_test_conf_param monitor_params[],
                     void **state)
 {
+    struct cli_protocol *prctx;
     errno_t ret;
 
     pam_test_ctx = talloc_zero(NULL, struct pam_test_ctx);
@@ -256,9 +257,12 @@ void test_pam_setup(struct sss_test_conf_param dom_params[],
     /* Create client context */
     pam_test_ctx->cctx = mock_cctx(pam_test_ctx, pam_test_ctx->rctx);
     assert_non_null(pam_test_ctx->cctx);
-
-    pam_test_ctx->cctx->cli_protocol_version = register_cli_protocol_version();
     pam_test_ctx->cctx->ev = pam_test_ctx->tctx->ev;
+
+    prctx = mock_prctx(pam_test_ctx->cctx);
+    assert_non_null(prctx);
+    pam_test_ctx->cctx->protocol_ctx = prctx;
+    prctx->cli_protocol_version = register_cli_protocol_version();
 }
 
 static void pam_test_setup_common(void)
@@ -418,11 +422,14 @@ void __real_sss_packet_get_body(struct sss_packet *packet,
 
 void __wrap_sss_cmd_done(struct cli_ctx *cctx, void *freectx)
 {
-    struct sss_packet *packet = cctx->creq->out;
+    struct cli_protocol *prctx;
+    struct sss_packet *packet;
     uint8_t *body;
     size_t blen;
     cmd_cb_fn_t check_cb;
 
+    prctx = talloc_get_type(cctx->protocol_ctx, struct cli_protocol);
+    packet = prctx->creq->out;
     assert_non_null(packet);
 
     check_cb = sss_mock_ptr_type(cmd_cb_fn_t);